import os, cv2
import random
import numpy as np
import pandas as pd
import mediapipe as mp

class Mask:
    def __init__(self):
        self.eye_indices_l = [130, 247, 30, 29, 27, 28, 56, 190, 243, 112, 26, 22, 23, 24, 110, 25]   # 16
        self.eye_indices_r = [463, 414, 286, 258, 257, 259, 260, 467, 359, 255, 339, 254, 253, 252, 256, 341]   # 16
        self.mouth_indices = [61, 185, 40, 39, 37, 0, 267, 269, 270, 409, 291, 375, 321, 405, 314, 17, 84, 181, 91, 146]   # 20
        fc_indices = [234, 162, 54, 67, 10, 297, 284, 389, 454, 361, 397, 379, 400, 152, 176, 150, 172, 132]  # 18
        self.total_indices = self.eye_indices_l + self.eye_indices_r + self.mouth_indices + fc_indices

        self.mp_face_mesh = mp.solutions.face_mesh.FaceMesh(
            static_image_mode=False,
            max_num_faces=1,
            min_detection_confidence=0.5)
        
    def get_maskimg(self):
        image_folder = 'image'
        image_filenames = [f for f in os.listdir(image_folder) if f.endswith(('.png'))]
        random_image_filename = random.choice(image_filenames)
        print('Change mask to  ' + random_image_filename)
        self.mask_img = cv2.cvtColor(cv2.imread(image_folder + '/' + random_image_filename, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2BGRA)
        image_basename = random_image_filename.split('.')[0]
        df = pd.read_csv(image_folder + '/labels_' + image_basename + '.csv', header=None)
        self.coordinates = df[[1, 2]].values

    def get_landmarks(self, results):
        face_landmarks = results.multi_face_landmarks[0].landmark

        face_points = []
        for landmark in face_landmarks:
            x = int(landmark.x * width)
            y = int(landmark.y * height)
            face_points.append([x, y])
        self.face_points_np = np.array(face_points) 

        eye_landmarks_l = [(int(face_landmarks[i].x * width), int(face_landmarks[i].y * height)) for i in self.eye_indices_l]
        eye_landmarks_r = [(int(face_landmarks[i].x * width), int(face_landmarks[i].y * height)) for i in self.eye_indices_r]
        mouth_landmarks = [(int(face_landmarks[i].x * width), int(face_landmarks[i].y * height)) for i in self.mouth_indices]
        total_landmarks = [(int(face_landmarks[i].x * width), int(face_landmarks[i].y * height)) for i in self.total_indices]
        
        self.eye_points_l = np.array(eye_landmarks_l, np.int32)
        self.eye_points_r = np.array(eye_landmarks_r, np.int32)
        self.mouth_points = np.array(mouth_landmarks, np.int32)
        self.total_points = np.array(total_landmarks, np.int32)

    def get_matrix_and_mask(self, frame):
        mask = np.zeros_like(frame, dtype=np.uint8)
        ym = 200 * np.ones(frame.shape, dtype=np.uint8)
        matrix, _ = cv2.findHomography(self.coordinates, self.total_points)
        transformed_img = cv2.warpPerspective(self.mask_img, matrix, (frame.shape[1], frame.shape[0]))

        # 使用凸包算法对关键点进行排序
        hull = cv2.convexHull(self.face_points_np.astype(np.int32))  # 返回凸包的顶点
        eye_hull_l = cv2.convexHull(self.eye_points_l.astype(np.int32))
        eye_hull_r = cv2.convexHull(self.eye_points_r.astype(np.int32))
        mouth_hull = cv2.convexHull(self.mouth_points.astype(np.int32))
        # 在遮罩图像上绘制人脸区域
        cv2.fillPoly(mask, [hull.astype(np.int32)], (255, 255, 255), cv2.LINE_AA)
        cv2.fillPoly(mask, [eye_hull_l.astype(np.int32)], (0, 0, 0, 255), cv2.LINE_AA)
        cv2.fillPoly(mask, [eye_hull_r.astype(np.int32)], (0, 0, 0, 255), cv2.LINE_AA)
        cv2.fillPoly(mask, [mouth_hull.astype(np.int32)], (0, 0, 0, 255), cv2.LINE_AA)

        # 将变换后的源图像贴到目标图像上
        transformed_img = cv2.cvtColor(transformed_img, cv2.COLOR_BGRA2BGR)
        result = cv2.bitwise_and(transformed_img, mask)
        cv2.imshow('test', result)
        result = cv2.add(result, cv2.bitwise_and(frame, 255 - mask))
        # result = cv2.seamlessClone(result, frame, ym, (320, 240), cv2.MIXED_CLONE)

        return result

if __name__ == '__main__':

    G_mask = Mask()
    G_mask.get_maskimg()

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print('Camera is not open!')
        exit()
    font = cv2.FONT_HERSHEY_COMPLEX_SMALL

    while True:
        ret, frame = cap.read()
        height, width, _ = frame.shape
        cv2.putText(frame, "Press 'F' to change mask.", (10, 20), font, 1, (0, 0, 255), 1)
        
        results = G_mask.mp_face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

        if results.multi_face_landmarks:
            G_mask.get_landmarks(results)
            result = G_mask.get_matrix_and_mask(frame)
            
            cv2.imshow('Result', result)
        else:
            cv2.imshow('Result', frame)

        keypressed = cv2.waitKey(1) & 0xFF
        if keypressed == ord('q'):
            break
        elif keypressed == ord('f'):
            try:
                G_mask.get_maskimg()
            except:
                continue

    cap.release()
    cv2.destroyAllWindows()
